#coding: utf-8
# Utility to read multuple capture files and merge into one CSV. This CSV will contains statistical information about these packets
# Usage: python analize_to_csv.py <OUTPUT FILE> <FILE 1> <FILE 2> ... <FILE N>

import sys
from scapy.all import rdpcap, ARP
from StringIO import StringIO
from math import log


def entropia(ips):
    N = sum(ips.values())
    Ps = [k/N for k in ips.values()]
    H = -sum([p * log(p, 2) for p in Ps])
    return H


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print "Usage: analize_to_csv.py <OUTPUT FILE> <FILE 1> <FILE 2> ... <FILE N>"
        sys.exit(1)

    requests = {}
    replays = {}
    ipssrc = {}
    ipsdst = {}

    def read_data():
        for arg in sys.argv[2:]:
            print "Reading:", arg

            result = StringIO()
            x = rdpcap(arg)
            for packet in x:
                if packet[ARP].op == 1:
                    src = packet[ARP].psrc
                    targ = packet[ARP].pdst
                    if not src in ipssrc:
                        ipssrc[src] = 0.0
                    if not src in requests:
                        requests[src] = {}
                    ipssrc[src]+= 1

                    if not targ in ipsdst:
                        ipsdst[targ] = 0.0
                    if not targ in requests[src]:
                        requests[src][targ] = 0

                    ipsdst[targ] += 1
                    requests[src][targ] += 1

                elif packet[ARP].op == 2:
                    src = packet[ARP].pdst
                    targ = packet[ARP].psrc

                    if not src in ipssrc:
                        ipssrc[src] = 0.0
                    if not src in replays:
                        replays[src] = 0

                    ipssrc[src] += 1
                    replays[src] += 1

                    if not targ in ipsdst:
                        ipsdst[targ] = 0.0
                    ipsdst[targ] += 1

    read_data()

    outputFile=open(sys.argv[1], 'w+')

    src_ip_count = len(ipssrc)
    dst_ip_count = len(ipsdst)
    src_packet_count = sum(ipssrc.values())
    dst_packet_count = sum(ipsdst.values())

    outputFile.write('--- Modelo Cátedra ---\n')
    outputFile.write('Entropia Src, Entropia Dst, Src Count, Dst Count, Paquetes Src, Paquetes DST, Total IPs \n')
    outputFile.write("%f, %f, %f, %f, %f, %f, %s\n" % (
        entropia(ipsdst), entropia(ipssrc),
        src_ip_count, dst_ip_count, src_packet_count, dst_packet_count,
        len(set(ipssrc.keys()) | set(ipsdst.keys()))))
    outputFile.write('\n\n')



    outputFile.write('IP, Src Count, Dst Count, P(Src), P(Dst)\n')
    for ip in set(ipssrc.keys()) | set(ipsdst.keys()):
        outputFile.write('%s, %s, %s, %f, %f \n' % (
            ip, ipssrc.get(ip, 0), ipsdst.get(ip, 0),
            float(ipssrc.get(ip, 0))/ src_packet_count,
            float(ipsdst.get(ip, 0))/ dst_packet_count
        ))

    outputFile.write('--- Modelo 1 ---\n')
    outputFile.write('Message Type, Count\n')
    requestsCount = 0
    for src in requests:
        for targ in requests[src]:
            requestsCount = requestsCount + requests[src][targ]
    outputFile.write('Requests, ' + str(requestsCount) + '\n')

    replaysCount = 0
    for src in replays:
        replaysCount = replaysCount + replays[src]
    outputFile.write('Replays, ' + str(replaysCount) + '\n')
    outputFile.write('\n\n')

    outputFile.write('--- Modelo 2 ---\n')
    outputFile.write('Request source, Count\n')

    for src in requests:
        count = 0
        for targ in requests[src]:
            count = count + requests[src][targ]
        outputFile.write(src + ', ' + str(count) + '\n')

    outputFile.write('\n\n')
    outputFile.write('--- Modelo 3 ---\n')
    outputFile.write('Replay source, Count\n')
    for src in replays:
        count = replays[src]
        outputFile.write(src + ', ' + str(count) + '\n')

    requestsHistogram={}

    for src in requests:
        for targ in requests[src]:
            if not targ in requestsHistogram:
                requestsHistogram[targ] = 0
            requestsHistogram[targ] = requestsHistogram[targ] + requests[src][targ]

    outputFile.write('\n\n')
    outputFile.write('--- Histogramas de Requests ---\n')
    outputFile.write('Target, Count\n')

    for targ in sorted(requestsHistogram):
        count = requestsHistogram[targ]
        outputFile.write(targ + ', ' + str(count) + '\n')

    outputFile.write('\n\n')
    outputFile.write('--- Grafo de Requests ---\n')
    outputFile.write('Source, Target, Count\n')

    for src in sorted(requests):
        for targ in requests[src]:
            outputFile.write(src + ', ' + targ + ', ' + str(requests[src][targ]) + '\n')

    outputFile.close()
